import numpy as np
import pandas as pd
import torch
import torch.nn as nn
from torch.utils.data import Dataset
import torch.nn.functional as F


class Excel_dataset(Dataset):

    def __init__(self, dir, if_normalize=False):
        super(Excel_dataset, self).__init__()

        data = pd.read_excel(dir, engine="openpyxl")

        nplist = data.T.to_numpy()
        data = nplist[1:-1].T
        self.data = np.float64(data)
        self.target = nplist[-1]

        # Tensor化
        self.data = np.array(self.data)
        self.data = torch.FloatTensor(self.data)
        self.target_num = np.array(self.target - 1)
        self.target_num = self.target_num.astype(float)
        self.target_num = torch.LongTensor(self.target_num)

        if if_normalize == True:
            self.data = F.normalize(self.data)

    def __getitem__(self, index):
        # tar = F.one_hot(self.target[index].to(torch.int64), len(self.target_type))
        # print(tar)

        return self.data[index], self.target_num[index]

    def __len__(self):
        return len(self.target)


def data_split(data, rate):
    train_l = int(len(data) * rate)
    test_l = len(data) - train_l
    """打乱数据集并且划分"""
    train_set, test_set = torch.utils.data.random_split(data, [train_l, test_l])
    return train_set, test_set



